MappingsConverter.java

/*******************************************************************************
 * Copyright (c) 2015 Jeff Martin.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the GNU Lesser General Public
 * License v3.0 which accompanies this distribution, and is available at
 * http://www.gnu.org/licenses/lgpl.html
 *
 * Contributors:
 * Jeff Martin - initial API and implementation
 ******************************************************************************/
package cuchaz.enigma.convert;

import java.util.*;
import java.util.jar.JarFile;

import com.google.common.collect.BiMap;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;
import com.google.common.collect.Sets;

import cuchaz.enigma.Deobfuscator;
import cuchaz.enigma.analysis.JarIndex;
import cuchaz.enigma.convert.ClassNamer.SidedClassNamer;
import cuchaz.enigma.mapping.*;

public class MappingsConverter
{
	
	public static ClassMatches computeClassMatches(JarFile sourceJar,
		JarFile destJar, Mappings mappings)
	{
		
		// index jars
		System.out.println("Indexing source jar...");
		JarIndex sourceIndex = new JarIndex();
		sourceIndex.indexJar(sourceJar, false);
		System.out.println("Indexing dest jar...");
		JarIndex destIndex = new JarIndex();
		destIndex.indexJar(destJar, false);
		
		// compute the matching
		ClassMatching matching =
			computeMatching(sourceJar, sourceIndex, destJar, destIndex, null);
		return new ClassMatches(matching.matches());
	}
	
	public static ClassMatching computeMatching(JarFile sourceJar,
		JarIndex sourceIndex, JarFile destJar, JarIndex destIndex,
		BiMap<ClassEntry, ClassEntry> knownMatches)
	{
		
		System.out.println("Iteratively matching classes");
		
		ClassMatching lastMatching = null;
		int round = 0;
		SidedClassNamer sourceNamer = null;
		SidedClassNamer destNamer = null;
		for(boolean useReferences : Arrays.asList(false, true))
		{
			
			int numUniqueMatchesLastTime = 0;
			if(lastMatching != null)
				numUniqueMatchesLastTime = lastMatching.uniqueMatches().size();
			
			while(true)
			{
				
				System.out.println("Round " + (++round) + "...");
				
				// init the matching with identity settings
				ClassMatching matching =
					new ClassMatching(new ClassIdentifier(sourceJar,
						sourceIndex, sourceNamer, useReferences),
						new ClassIdentifier(destJar, destIndex, destNamer,
							useReferences));
				
				if(knownMatches != null)
					matching.addKnownMatches(knownMatches);
				
				if(lastMatching == null)
					// search all classes
					matching.match(sourceIndex.getObfClassEntries(),
						destIndex.getObfClassEntries());
				else
				{
					// we already know about these matches from last time
					matching.addKnownMatches(lastMatching.uniqueMatches());
					
					// search unmatched and ambiguously-matched classes
					matching.match(lastMatching.unmatchedSourceClasses(),
						lastMatching.unmatchedDestClasses());
					for(ClassMatch match : lastMatching.ambiguousMatches())
						matching.match(match.sourceClasses, match.destClasses);
				}
				System.out.println(matching);
				BiMap<ClassEntry, ClassEntry> uniqueMatches =
					matching.uniqueMatches();
				
				// did we match anything new this time?
				if(uniqueMatches.size() > numUniqueMatchesLastTime)
				{
					numUniqueMatchesLastTime = uniqueMatches.size();
					lastMatching = matching;
				}else
					break;
				
				// update the namers
				ClassNamer namer = new ClassNamer(uniqueMatches);
				sourceNamer = namer.getSourceNamer();
				destNamer = namer.getDestNamer();
			}
		}
		
		return lastMatching;
	}
	
	public static Mappings newMappings(ClassMatches matches,
		Mappings oldMappings, Deobfuscator sourceDeobfuscator,
		Deobfuscator destDeobfuscator)
	{
		
		// sort the unique matches by size of inner class chain
		Multimap<Integer, java.util.Map.Entry<ClassEntry, ClassEntry>> matchesByDestChainSize =
			HashMultimap.create();
		for(java.util.Map.Entry<ClassEntry, ClassEntry> match : matches
			.getUniqueMatches().entrySet())
		{
			int chainSize =
				destDeobfuscator.getJarIndex()
					.getObfClassChain(match.getValue()).size();
			matchesByDestChainSize.put(chainSize, match);
		}
		
		// build the mappings (in order of small-to-large inner chains)
		Mappings newMappings = new Mappings();
		List<Integer> chainSizes =
			Lists.newArrayList(matchesByDestChainSize.keySet());
		Collections.sort(chainSizes);
		for(int chainSize : chainSizes)
			for(java.util.Map.Entry<ClassEntry, ClassEntry> match : matchesByDestChainSize
				.get(chainSize))
			{
				
				// get class info
				ClassEntry obfSourceClassEntry = match.getKey();
				ClassEntry obfDestClassEntry = match.getValue();
				List<ClassEntry> destClassChain =
					destDeobfuscator.getJarIndex().getObfClassChain(
						obfDestClassEntry);
				
				ClassMapping sourceMapping =
					sourceDeobfuscator.getMappings().getClassByObf(
						obfSourceClassEntry);
				if(sourceMapping == null)
					// if this class was never deobfuscated, don't try to match
					// it
					continue;
				
				// find out where to make the dest class mapping
				if(destClassChain.size() == 1)
					// not an inner class, add directly to mappings
					newMappings.addClassMapping(migrateClassMapping(
						obfDestClassEntry, sourceMapping, matches, false));
				else
				{
					// inner class, find the outer class mapping
					ClassMapping destMapping = null;
					for(int i = 0; i < destClassChain.size() - 1; i++)
					{
						ClassEntry destChainClassEntry = destClassChain.get(i);
						if(destMapping == null)
						{
							destMapping =
								newMappings.getClassByObf(destChainClassEntry);
							if(destMapping == null)
							{
								destMapping =
									new ClassMapping(
										destChainClassEntry.getName());
								newMappings.addClassMapping(destMapping);
							}
						}else
						{
							destMapping =
								destMapping
									.getInnerClassByObfSimple(destChainClassEntry
										.getInnermostClassName());
							if(destMapping == null)
							{
								destMapping =
									new ClassMapping(
										destChainClassEntry.getName());
								destMapping.addInnerClassMapping(destMapping);
							}
						}
					}
					destMapping.addInnerClassMapping(migrateClassMapping(
						obfDestClassEntry, sourceMapping, matches, true));
				}
			}
		return newMappings;
	}
	
	private static ClassMapping migrateClassMapping(ClassEntry newObfClass,
		ClassMapping oldClassMapping, final ClassMatches matches,
		boolean useSimpleName)
	{
		
		ClassNameReplacer replacer = new ClassNameReplacer()
		{
			@Override
			public String replace(String className)
			{
				ClassEntry newClassEntry =
					matches.getUniqueMatches().get(new ClassEntry(className));
				if(newClassEntry != null)
					return newClassEntry.getName();
				return null;
			}
		};
		
		ClassMapping newClassMapping;
		String deobfName = oldClassMapping.getDeobfName();
		if(deobfName != null)
		{
			if(useSimpleName)
				deobfName = new ClassEntry(deobfName).getSimpleName();
			newClassMapping =
				new ClassMapping(newObfClass.getName(), deobfName);
		}else
			newClassMapping = new ClassMapping(newObfClass.getName());
		
		// copy fields
		for(FieldMapping fieldMapping : oldClassMapping.fields())
			newClassMapping.addFieldMapping(new FieldMapping(fieldMapping,
				replacer));
		
		// copy methods
		for(MethodMapping methodMapping : oldClassMapping.methods())
			newClassMapping.addMethodMapping(new MethodMapping(methodMapping,
				replacer));
		
		return newClassMapping;
	}
	
	public static void convertMappings(Mappings mappings,
		BiMap<ClassEntry, ClassEntry> changes)
	{
		
		// sort the changes so classes are renamed in the correct order
		// ie. if we have the mappings a->b, b->c, we have to apply b->c before
		// a->b
		LinkedHashMap<ClassEntry, ClassEntry> sortedChanges =
			Maps.newLinkedHashMap();
		int numChangesLeft = changes.size();
		while(!changes.isEmpty())
		{
			Iterator<Map.Entry<ClassEntry, ClassEntry>> iter =
				changes.entrySet().iterator();
			while(iter.hasNext())
			{
				Map.Entry<ClassEntry, ClassEntry> change = iter.next();
				if(changes.containsKey(change.getValue()))
				{
					sortedChanges.put(change.getKey(), change.getValue());
					iter.remove();
				}
			}
			
			// did we remove any changes?
			if(numChangesLeft - changes.size() > 0)
				// keep going
				numChangesLeft = changes.size();
			else
				// can't sort anymore. There must be a loop
				break;
		}
		if(!changes.isEmpty())
			throw new Error(
				"Unable to sort class changes! There must be a cycle.");
		
		// convert the mappings in the correct class order
		for(Map.Entry<ClassEntry, ClassEntry> entry : sortedChanges.entrySet())
			mappings.renameObfClass(entry.getKey().getName(), entry.getValue()
				.getName());
	}
	
	public static interface Doer<T extends Entry>
	{
		Collection<T> getDroppedEntries(MappingsChecker checker);
		
		Collection<T> getObfEntries(JarIndex jarIndex);
		
		Collection<? extends MemberMapping<T>> getMappings(
			ClassMapping destClassMapping);
		
		Set<T> filterEntries(Collection<T> obfEntries, T obfSourceEntry,
			ClassMatches classMatches);
		
		void setUpdateObfMember(ClassMapping classMapping,
			MemberMapping<T> memberMapping, T newEntry);
		
		boolean hasObfMember(ClassMapping classMapping, T obfEntry);
		
		void removeMemberByObf(ClassMapping classMapping, T obfEntry);
	}
	
	public static Doer<FieldEntry> getFieldDoer()
	{
		return new Doer<FieldEntry>()
		{
			
			@Override
			public Collection<FieldEntry> getDroppedEntries(
				MappingsChecker checker)
			{
				return checker.getDroppedFieldMappings().keySet();
			}
			
			@Override
			public Collection<FieldEntry> getObfEntries(JarIndex jarIndex)
			{
				return jarIndex.getObfFieldEntries();
			}
			
			@Override
			public Collection<? extends MemberMapping<FieldEntry>> getMappings(
				ClassMapping destClassMapping)
			{
				return (Collection<? extends MemberMapping<FieldEntry>>)destClassMapping
					.fields();
			}
			
			@Override
			public Set<FieldEntry> filterEntries(
				Collection<FieldEntry> obfDestFields,
				FieldEntry obfSourceField, ClassMatches classMatches)
			{
				Set<FieldEntry> out = Sets.newHashSet();
				for(FieldEntry obfDestField : obfDestFields)
				{
					Type translatedDestType =
						translate(obfDestField.getType(), classMatches
							.getUniqueMatches().inverse());
					if(translatedDestType.equals(obfSourceField.getType()))
						out.add(obfDestField);
				}
				return out;
			}
			
			@Override
			public void setUpdateObfMember(ClassMapping classMapping,
				MemberMapping<FieldEntry> memberMapping, FieldEntry newField)
			{
				FieldMapping fieldMapping = (FieldMapping)memberMapping;
				classMapping.setFieldObfNameAndType(fieldMapping.getObfName(),
					fieldMapping.getObfType(), newField.getName(),
					newField.getType());
			}
			
			@Override
			public boolean hasObfMember(ClassMapping classMapping,
				FieldEntry obfField)
			{
				return classMapping.getFieldByObf(obfField.getName(),
					obfField.getType()) != null;
			}
			
			@Override
			public void removeMemberByObf(ClassMapping classMapping,
				FieldEntry obfField)
			{
				classMapping.removeFieldMapping(classMapping.getFieldByObf(
					obfField.getName(), obfField.getType()));
			}
		};
	}
	
	public static Doer<BehaviorEntry> getMethodDoer()
	{
		return new Doer<BehaviorEntry>()
		{
			
			@Override
			public Collection<BehaviorEntry> getDroppedEntries(
				MappingsChecker checker)
			{
				return checker.getDroppedMethodMappings().keySet();
			}
			
			@Override
			public Collection<BehaviorEntry> getObfEntries(JarIndex jarIndex)
			{
				return jarIndex.getObfBehaviorEntries();
			}
			
			@Override
			public Collection<? extends MemberMapping<BehaviorEntry>> getMappings(
				ClassMapping destClassMapping)
			{
				return (Collection<? extends MemberMapping<BehaviorEntry>>)destClassMapping
					.methods();
			}
			
			@Override
			public Set<BehaviorEntry> filterEntries(
				Collection<BehaviorEntry> obfDestFields,
				BehaviorEntry obfSourceField, ClassMatches classMatches)
			{
				Set<BehaviorEntry> out = Sets.newHashSet();
				for(BehaviorEntry obfDestField : obfDestFields)
				{
					Signature translatedDestSignature =
						translate(obfDestField.getSignature(), classMatches
							.getUniqueMatches().inverse());
					if(translatedDestSignature == null
						&& obfSourceField.getSignature() == null)
						out.add(obfDestField);
					else if(translatedDestSignature == null
						|| obfSourceField.getSignature() == null)
					{
						// skip it
					}else if(translatedDestSignature.equals(obfSourceField
						.getSignature()))
						out.add(obfDestField);
				}
				return out;
			}
			
			@Override
			public void setUpdateObfMember(ClassMapping classMapping,
				MemberMapping<BehaviorEntry> memberMapping,
				BehaviorEntry newBehavior)
			{
				MethodMapping methodMapping = (MethodMapping)memberMapping;
				classMapping.setMethodObfNameAndSignature(
					methodMapping.getObfName(),
					methodMapping.getObfSignature(), newBehavior.getName(),
					newBehavior.getSignature());
			}
			
			@Override
			public boolean hasObfMember(ClassMapping classMapping,
				BehaviorEntry obfBehavior)
			{
				return classMapping.getMethodByObf(obfBehavior.getName(),
					obfBehavior.getSignature()) != null;
			}
			
			@Override
			public void removeMemberByObf(ClassMapping classMapping,
				BehaviorEntry obfBehavior)
			{
				classMapping.removeMethodMapping(classMapping.getMethodByObf(
					obfBehavior.getName(), obfBehavior.getSignature()));
			}
		};
	}
	
	public static <T extends Entry> MemberMatches<T> computeMemberMatches(
		Deobfuscator destDeobfuscator, Mappings destMappings,
		ClassMatches classMatches, Doer<T> doer)
	{
		
		MemberMatches<T> memberMatches = new MemberMatches<T>();
		
		// unmatched source fields are easy
		MappingsChecker checker =
			new MappingsChecker(destDeobfuscator.getJarIndex());
		checker.dropBrokenMappings(destMappings);
		for(T destObfEntry : doer.getDroppedEntries(checker))
		{
			T srcObfEntry =
				translate(destObfEntry, classMatches.getUniqueMatches()
					.inverse());
			memberMatches.addUnmatchedSourceEntry(srcObfEntry);
		}
		
		// get matched fields (anything that's left after the checks/drops is
		// matched(
		for(ClassMapping classMapping : destMappings.classes())
			collectMatchedFields(memberMatches, classMapping, classMatches,
				doer);
		
		// get unmatched dest fields
		for(T destEntry : doer.getObfEntries(destDeobfuscator.getJarIndex()))
			if(!memberMatches.isMatchedDestEntry(destEntry))
				memberMatches.addUnmatchedDestEntry(destEntry);
		
		System.out.println("Automatching "
			+ memberMatches.getUnmatchedSourceEntries().size()
			+ " unmatched source entries...");
		
		// go through the unmatched source fields and try to pick out the easy
		// matches
		for(ClassEntry obfSourceClass : Lists.newArrayList(memberMatches
			.getSourceClassesWithUnmatchedEntries()))
			for(T obfSourceEntry : Lists.newArrayList(memberMatches
				.getUnmatchedSourceEntries(obfSourceClass)))
			{
				
				// get the possible dest matches
				ClassEntry obfDestClass =
					classMatches.getUniqueMatches().get(obfSourceClass);
				
				// filter by type/signature
				Set<T> obfDestEntries =
					doer.filterEntries(
						memberMatches.getUnmatchedDestEntries(obfDestClass),
						obfSourceEntry, classMatches);
				
				if(obfDestEntries.size() == 1)
					// make the easy match
					memberMatches.makeMatch(obfSourceEntry, obfDestEntries
						.iterator().next());
				else if(obfDestEntries.isEmpty())
					// no match is possible =(
					memberMatches.makeSourceUnmatchable(obfSourceEntry);
			}
		
		System.out.println(String.format(
			"Ended up with %d ambiguous and %d unmatchable source entries",
			memberMatches.getUnmatchedSourceEntries().size(), memberMatches
				.getUnmatchableSourceEntries().size()));
		
		return memberMatches;
	}
	
	private static <T extends Entry> void collectMatchedFields(
		MemberMatches<T> memberMatches, ClassMapping destClassMapping,
		ClassMatches classMatches, Doer<T> doer)
	{
		
		// get the fields for this class
		for(MemberMapping<T> destEntryMapping : doer
			.getMappings(destClassMapping))
		{
			T destObfField =
				destEntryMapping.getObfEntry(destClassMapping.getObfEntry());
			T srcObfField =
				translate(destObfField, classMatches.getUniqueMatches()
					.inverse());
			memberMatches.addMatch(srcObfField, destObfField);
		}
		
		// recurse
		for(ClassMapping destInnerClassMapping : destClassMapping
			.innerClasses())
			collectMatchedFields(memberMatches, destInnerClassMapping,
				classMatches, doer);
	}
	
	@SuppressWarnings("unchecked")
	private static <T extends Entry> T translate(T in,
		BiMap<ClassEntry, ClassEntry> map)
	{
		if(in instanceof FieldEntry)
			return (T)new FieldEntry(map.get(in.getClassEntry()), in.getName(),
				translate(((FieldEntry)in).getType(), map));
		else if(in instanceof MethodEntry)
			return (T)new MethodEntry(map.get(in.getClassEntry()),
				in.getName(), translate(((MethodEntry)in).getSignature(), map));
		else if(in instanceof ConstructorEntry)
			return (T)new ConstructorEntry(map.get(in.getClassEntry()),
				translate(((ConstructorEntry)in).getSignature(), map));
		throw new Error("Unhandled entry type: " + in.getClass());
	}
	
	private static Type translate(Type type,
		final BiMap<ClassEntry, ClassEntry> map)
	{
		return new Type(type, new ClassNameReplacer()
		{
			@Override
			public String replace(String inClassName)
			{
				ClassEntry outClassEntry = map.get(new ClassEntry(inClassName));
				if(outClassEntry == null)
					return null;
				return outClassEntry.getName();
			}
		});
	}
	
	private static Signature translate(Signature signature,
		final BiMap<ClassEntry, ClassEntry> map)
	{
		if(signature == null)
			return null;
		return new Signature(signature, new ClassNameReplacer()
		{
			@Override
			public String replace(String inClassName)
			{
				ClassEntry outClassEntry = map.get(new ClassEntry(inClassName));
				if(outClassEntry == null)
					return null;
				return outClassEntry.getName();
			}
		});
	}
	
	public static <T extends Entry> void applyMemberMatches(Mappings mappings,
		ClassMatches classMatches, MemberMatches<T> memberMatches, Doer<T> doer)
	{
		for(ClassMapping classMapping : mappings.classes())
			applyMemberMatches(classMapping, classMatches, memberMatches, doer);
	}
	
	private static <T extends Entry> void applyMemberMatches(
		ClassMapping classMapping, ClassMatches classMatches,
		MemberMatches<T> memberMatches, Doer<T> doer)
	{
		
		// get the classes
		ClassEntry obfDestClass = classMapping.getObfEntry();
		
		// make a map of all the renames we need to make
		Map<T, T> renames = Maps.newHashMap();
		for(MemberMapping<T> memberMapping : Lists.newArrayList(doer
			.getMappings(classMapping)))
		{
			T obfOldDestEntry = memberMapping.getObfEntry(obfDestClass);
			T obfSourceEntry =
				getSourceEntryFromDestMapping(memberMapping, obfDestClass,
					classMatches);
			
			// but drop the unmatchable things
			if(memberMatches.isUnmatchableSourceEntry(obfSourceEntry))
			{
				doer.removeMemberByObf(classMapping, obfOldDestEntry);
				continue;
			}
			
			T obfNewDestEntry = memberMatches.matches().get(obfSourceEntry);
			if(obfNewDestEntry != null
				&& !obfOldDestEntry.getName().equals(obfNewDestEntry.getName()))
				renames.put(obfOldDestEntry, obfNewDestEntry);
		}
		
		if(!renames.isEmpty())
		{
			
			// apply to this class (should never need more than n passes)
			int numRenamesAppliedThisRound;
			do
			{
				numRenamesAppliedThisRound = 0;
				
				for(MemberMapping<T> memberMapping : Lists.newArrayList(doer
					.getMappings(classMapping)))
				{
					T obfOldDestEntry = memberMapping.getObfEntry(obfDestClass);
					T obfNewDestEntry = renames.get(obfOldDestEntry);
					if(obfNewDestEntry != null)
						// make sure this rename won't cause a collision
						// otherwise, save it for the next round and try again
						// next time
						if(!doer.hasObfMember(classMapping, obfNewDestEntry))
						{
							doer.setUpdateObfMember(classMapping,
								memberMapping, obfNewDestEntry);
							renames.remove(obfOldDestEntry);
							numRenamesAppliedThisRound++;
						}
				}
			}while(numRenamesAppliedThisRound > 0);
			
			if(!renames.isEmpty())
			{
				System.err
					.println(String
						.format(
							"WARNING: Couldn't apply all the renames for class %s. %d renames left.",
							classMapping.getObfFullName(), renames.size()));
				for(Map.Entry<T, T> entry : renames.entrySet())
					System.err.println(String.format("\t%s -> %s", entry
						.getKey().getName(), entry.getValue().getName()));
			}
		}
		
		// recurse
		for(ClassMapping innerClassMapping : classMapping.innerClasses())
			applyMemberMatches(innerClassMapping, classMatches, memberMatches,
				doer);
	}
	
	private static <T extends Entry> T getSourceEntryFromDestMapping(
		MemberMapping<T> destMemberMapping, ClassEntry obfDestClass,
		ClassMatches classMatches)
	{
		return translate(destMemberMapping.getObfEntry(obfDestClass),
			classMatches.getUniqueMatches().inverse());
	}
}